// copyright:no copyright,use it if you like
// license boilerplate:BSD 3-Clause
// filename:Multigrid.hpp
// author:Tempest(https://gitee.com/Tempest56890)
// date:2023-03-31
// brief:implement the Multigrid Method.

#ifndef __TEMPEST_MULTIGRID__
#define __TEMPEST_MULTIGRID__

#include "Matrix.h"
#include "Function.h"
#include <vector>

#define pi 3.1415926536

//构建模版类，以dim为参数，分别求解一维、二维问题.
template<int dim>
class MULTIGRID {};

//求解一维问题，问题域为(0,1).
template <>
class MULTIGRID<1> {
private:
    //n描述最粗网格的网格数，cycle_depth描述最细网格数之于最粗网格数的倍数
    int n,cycle_depth;
    //h描述最细网格的网格间距
    double h;
    //x存储最细网格的网格点(包含0和1两个点)
    std::vector<double> x;
    //f存储最细网格上的边界条件;u存储求得的函数值,不妨在边界点上为0
    std::vector<double> f,u;
    //type存储边界条件类型
    std::string type;
    //u_func为所求函数，仅用于消除常数差以及测试时的误差分析，输入时保证边界点上为0;func存储条件函数
    Function_one_dim &u_func,&func;
    Matrix<double> V__cycle(Matrix<double> v_h,Matrix<double> f_h,int n1=256,int n2=256,int nn=3) {
        Matrix<double> Tw_h(nn-1,nn-1),A_h(nn-1,nn-1),f__h(nn-1,1),res(nn-1,1);
        Tw_h.makeEmpty();A_h.makeEmpty();res.makeEmpty();
        for(int i=0;i<nn-1;i++) {
            res.reset(i,0,v_h(i,0));
        }
        for(int i=0;i<nn-2;i++) {
            Tw_h.reset(i,i,1.0/3.0);A_h.reset(i,i,2*nn*nn);
            Tw_h.reset(i,i+1,1.0/3.0);A_h.reset(i,i+1,-nn*nn);
            Tw_h.reset(i+1,i,1.0/3.0);A_h.reset(i+1,i,-nn*nn);
        }
        Tw_h.reset(nn-2,nn-2,1.0/3.0);A_h.reset(nn-2,nn-2,2*nn*nn);
        for(int i=0;i<nn-1;i++) f__h.reset(i,0,f_h(i,0)/(3.0*nn*nn));
        for(int iter=1;iter<=n1;iter++){
            res=Tw_h*res;
            for(int i=0;i<nn-1;i++) res.reset(i,0,res(i,0)+f__h(i,0));
        }
        if(nn!=n) {
            Matrix<double> f_2h=full_weighting(f_h-A_h*res,nn);
            Matrix<double> v_2h(nn/2-1,1);v_2h.makeEmpty();
            v_2h=V__cycle(v_2h,f_2h,n1,n2,nn/2);
            res=res+linear_interpolation(v_2h,nn/2);
            for(int iter=1;iter<=n2;iter++){
                res=Tw_h*res;
                for(int i=0;i<nn-1;i++) res.reset(i,0,res(i,0)+f__h(i,0));
            }
        } else {
            for(int iter=1;iter<=n2;iter++){
                res=Tw_h*res;
                for(int i=0;i<nn-1;i++) res.reset(i,0,res(i,0)+f__h(i,0));
            }
        }
        return res;
    }
    Matrix<double> FMG(Matrix<double> f_h,int n1=256,int n2=256,int nn=3) {
        if(nn!=n) {
            Matrix<double> f_2h=full_weighting(f_h,nn);
            Matrix<double> v_2h=FMG(f_2h,n1,n2,nn/2);
            Matrix<double> v_h=linear_interpolation(v_2h,nn/2);
            v_h=V__cycle(v_h,f_h,n1,n2,nn);
            return v_h;
        } else {
            Matrix<double> v_h(nn-1,1);
            v_h.makeEmpty();
            v_h=V__cycle(v_h,f_h,n1,n2,nn);
            return v_h;
        }
    }
    Matrix<double> injection(Matrix<double> vec,int nnn) {
        Matrix<double> res(nnn/2-1,1);
        res.makeEmpty();
        for(int i=0;i<nnn/2-1;i++) {
            res.reset(i,0,vec(2*i+1,0));
        }
        return res;
    }
    Matrix<double> full_weighting(Matrix<double> vec,int nnn) {
        Matrix<double> res(nnn/2-1,1);
        res.makeEmpty();
        for(int i=0;i<nnn/2-1;i++) {
            res.reset(i,0,vec(2*i+1,0)/2.0+vec(2*i,0)/4.0+vec(2*i+2,0)/4.0);
        }
        return res;
    }
    Matrix<double> linear_interpolation(Matrix<double> vec,int nnn){
        Matrix<double> I_2h_h(2*nnn-1,nnn-1),res(2*nnn-1,1);
        res.makeEmpty();
        for(int i=0;i<nnn-1;i++){
            I_2h_h.reset(2*i,i,1.0/2);
            I_2h_h.reset(2*i+1,i,1.0);
            I_2h_h.reset(2*i+2,i,1.0/2);
        }
        res=I_2h_h*vec;
        return res;
    }
public:
    MULTIGRID(Function_one_dim &u_func,Function_one_dim &func,int n=3,int cycle_depth=2,std::string type="O")
    :u_func(u_func),func(func),n(n),cycle_depth(cycle_depth),type(type){
        h=1.0/(n*pow(2,cycle_depth));
        x.clear();f.clear();u.clear();
        for(int i=0;i<n*pow(2,cycle_depth)+1;i++) {
            x.push_back(i*h);
            f.push_back(func(i*h));
            u.push_back(0.0);
        }
    }
    void V_cycle() {
        Matrix<double> v0(n*pow(2,cycle_depth)-1,1),f0(n*pow(2,cycle_depth)-1,1);
        v0.makeEmpty();
        for(int i=1;i<n*pow(2,cycle_depth);i++) f0.reset(i-1,0,f[i]);
        Matrix<double> res=V__cycle(v0,f0,256,256,n*pow(2,cycle_depth));
        double sum=0.0;
        for(int i=1;i<n*pow(2,cycle_depth);i++){
            sum+=u_func(x[i])-res(i-1,0);
        }
        sum=sum/(n*pow(2,cycle_depth)-1);
        for(int i=1;i<n*pow(2,cycle_depth);i++){
            u[i]=res(i-1,0)+sum;
        }
    }
    void FMG() {
        Matrix<double> f0(n*pow(2,cycle_depth)-1,1);
        for(int i=1;i<n*pow(2,cycle_depth);i++) f0.reset(i-1,0,f[i]);
        Matrix<double> v0=FMG(f0,256,256,n*pow(2,cycle_depth));
        double sum=0.0;
        for(int i=1;i<n*pow(2,cycle_depth);i++){
            sum+=u_func(x[i])-v0(i-1,0);
        }
        sum=sum/(n*pow(2,cycle_depth)-1);
        for(int i=1;i<n*pow(2,cycle_depth);i++){
            u[i]=v0(i-1,0)+sum;
        }
    }
    void relative_error_print() {
        for(int i=0;i<n*pow(2,cycle_depth)+1;i++) {
            std::cout<<(u_func(x[i])-u[i])/u_func(x[i])<<", ";
        }
    }
    void absolute_error_2norm_print(int j) {
        double sum=0.0;
        for(int i=0;i<n*pow(2,cycle_depth)+1;i++) {
            sum+=pow(u_func(x[i])-u[i],2.0);
        }
        sum=pow(sum/(n*pow(2,cycle_depth)),0.5);
        std::cout<<"test"<<j<<"'s "<<"absolute error in 2-norm: "<<sum<<" "<<std::endl;
    }
    void relative_error_2norm_print(int j) {
        double sum1=0.0,sum2=0.0;
        for(int i=0;i<n*pow(2,cycle_depth)+1;i++) {
            sum1+=pow(u_func(x[i])-u[i],2.0);
            sum2+=pow(u_func(x[i]),2.0);
        }
        sum1=pow(sum1/(n*pow(2,cycle_depth)),0.5);
        sum2=pow(sum2/(n*pow(2,cycle_depth)),0.5);
        std::cout<<"test"<<j<<"'s "<<"relative error in 2-norm: "<<sum1/sum2<<" "<<std::endl;
    }
};

//求解二维问题，问题域为(0,1)*(0,1).
template<>
class MULTIGRID<2> {
private:
    //n描述最粗网格的网格数(横向)，cycle_depth描述最细网格数(横向)之于最粗网格数(横向)的倍数
    int n,cycle_depth;
    //h描述最细网格的网格间距
    double h;
    //(x,y)存储最细网格的网格点(包含(0,0),(0,1),(1,0),(1,1)等点)
    std::vector<double> x,y;
    //f存储最细网格上的边界条件;u存储求得的函数值,不妨在边界点上为0
    std::vector<double> f,u;
    //type存储边界条件类型
    std::string type;
    //u_func为所求函数，仅用于消除常数差以及测试时的误差分析，输入时保证边界点上为0;func存储条件函数
    Function_two_dim &u_func,&func;
    
    Matrix<double> V__cycle(Matrix<double> v_h,Matrix<double> f_h,int n1=256,int n2=256,int nn=3) {
        Matrix<double> Tw_h(nn-1,nn-1),A_h(nn-1,nn-1),f__h(nn-1,1),res(nn-1,1);
        Tw_h.makeEmpty();A_h.makeEmpty();res.makeEmpty();
        for(int i=0;i<nn-1;i++) {
            res.reset(i,0,v_h(i,0));
        }
        for(int i=0;i<nn-2;i++) {
            Tw_h.reset(i,i,1.0/3.0);A_h.reset(i,i,2*nn*nn);
            Tw_h.reset(i,i+1,1.0/3.0);A_h.reset(i,i+1,-nn*nn);
            Tw_h.reset(i+1,i,1.0/3.0);A_h.reset(i+1,i,-nn*nn);
        }
        Tw_h.reset(nn-2,nn-2,1.0/3.0);A_h.reset(nn-2,nn-2,2*nn*nn);
        for(int i=0;i<nn-1;i++) f__h.reset(i,0,f_h(i,0)/(3.0*nn*nn));
        for(int iter=1;iter<=n1;iter++){
            res=Tw_h*res;
            for(int i=0;i<nn-1;i++) res.reset(i,0,res(i,0)+f__h(i,0));
        }
        if(nn!=n) {
            Matrix<double> f_2h=full_weighting(f_h-A_h*res,nn);
            Matrix<double> v_2h(nn/2-1,1);v_2h.makeEmpty();
            v_2h=V__cycle(v_2h,f_2h,n1,n2,nn/2);
            res=res+linear_interpolation(v_2h,nn/2);
            for(int iter=1;iter<=n2;iter++){
                res=Tw_h*res;
                for(int i=0;i<nn-1;i++) res.reset(i,0,res(i,0)+f__h(i,0));
            }
        } else {
            for(int iter=1;iter<=n2;iter++){
                res=Tw_h*res;
                for(int i=0;i<nn-1;i++) res.reset(i,0,res(i,0)+f__h(i,0));
            }
        }
        return res;
    }
    Matrix<double> FMG(Matrix<double> f_h,int n1=256,int n2=256,int nn=3) {
        if(nn!=n) {
            Matrix<double> f_2h=full_weighting(f_h,nn);
            Matrix<double> v_2h=FMG(f_2h,n1,n2,nn/2);
            Matrix<double> v_h=linear_interpolation(v_2h,nn/2);
            v_h=V__cycle(v_h,f_h,n1,n2,nn);
            return v_h;
        } else {
            Matrix<double> v_h(nn-1,1);
            v_h.makeEmpty();
            v_h=V__cycle(v_h,f_h,n1,n2,nn);
            return v_h;
        }
    }
    Matrix<double> injection(Matrix<double> vec,int nnn) {
        Matrix<double> res(nnn/2-1,1);
        res.makeEmpty();
        for(int i=0;i<nnn/2-1;i++) {
            res.reset(i,0,vec(2*i+1,0));
        }
        return res;
    }
    Matrix<double> full_weighting(Matrix<double> vec,int nnn) {
        Matrix<double> res(nnn/2-1,1);
        res.makeEmpty();
        for(int i=0;i<nnn/2-1;i++) {
            res.reset(i,0,vec(2*i+1,0)/2.0+vec(2*i,0)/4.0+vec(2*i+2,0)/4.0);
        }
        return res;
    }
    Matrix<double> linear_interpolation(Matrix<double> vec,int nnn){
        Matrix<double> I_2h_h(2*nnn-1,nnn-1),res(2*nnn-1,1);
        res.makeEmpty();
        for(int i=0;i<nnn-1;i++){
            I_2h_h.reset(2*i,i,1.0/2);
            I_2h_h.reset(2*i+1,i,1.0);
            I_2h_h.reset(2*i+2,i,1.0/2);
        }
        res=I_2h_h*vec;
        return res;
    }
public:
    MULTIGRID(Function_two_dim &u_func,Function_two_dim &func,int n=3,int cycle_depth=2,std::string type="O")
    :u_func(u_func),func(func),n(n),cycle_depth(cycle_depth),type(type){
        h=1.0/(n*pow(2,cycle_depth));
        x.clear();y.clear();f.clear();u.clear();
        for(int i=0;i<n*pow(2,cycle_depth)+1;i++) {
            x.push_back(i*h);
            y.push_back(i*h);
            for(int j=0;j<n*pow(2,cycle_depth)+1;j++) {
                u.push_back(0.0);
                f.push_back(func(j*h,i*h));
            }
        }
    }
    void V_cycle() {
        for(int j=0;j<=n*pow(2,cycle_depth);j++) {
            Matrix<double> v0(n*pow(2,cycle_depth)-1,1),f0(n*pow(2,cycle_depth)-1,1);
            v0.makeEmpty();
            for(int i=1;i<n*pow(2,cycle_depth);i++)
                f0.reset(i-1,0,f[j*(n*pow(2,cycle_depth)+1)+i]);
            Matrix<double> res=V__cycle(v0,f0,256,256,n*pow(2,cycle_depth));
            double sum=0.0;
            for(int i=1;i<n*pow(2,cycle_depth);i++){
                sum+=u_func(x[i],y[j])-res(i-1,0);
            }
            sum=sum/(n*pow(2,cycle_depth)-1);
            for(int i=1;i<n*pow(2,cycle_depth);i++){
                u[j*(n*pow(2,cycle_depth)+1)+i]=res(i-1,0)+sum;
            }
        }
    }
    void FMG() {
        for(int j=0;j<=n*pow(2,cycle_depth);j++) {
            Matrix<double> f0(n*pow(2,cycle_depth)-1,1);
            for(int i=1;i<n*pow(2,cycle_depth);i++)
                f0.reset(i-1,0,f[j*(n*pow(2,cycle_depth)+1)+i]);
            Matrix<double> v0=FMG(f0,256,256,n*pow(2,cycle_depth));
            double sum=0.0;
            for(int i=1;i<n*pow(2,cycle_depth);i++){
                sum+=u_func(x[i],y[j])-v0(i-1,0);
            }
            sum=sum/(n*pow(2,cycle_depth)-1);
            for(int i=1;i<n*pow(2,cycle_depth);i++){
                u[j*(n*pow(2,cycle_depth)+1)+i]=v0(i-1,0)+sum;
            }
        }
    }
    void relative_error_print() {
        for(int i=0;i<(n*pow(2,cycle_depth)+1);i++) {
            for(int j=0;j<(n*pow(2,cycle_depth)+1);j++) {
                std::cout<<(u_func(x[j],y[i])-u[i*(n*pow(2,cycle_depth)+1)+j])<<", ";
            }//(u_func(x[j],y[i])-u[i*(n*pow(2,cycle_depth)+1)+j])
        }
    }
    void absolute_error_2norm_print(int j) {
        double sum=0.0;
        for(int i=0;i<n*pow(2,cycle_depth)+1;i++) {
            for(int j=0;j<n*pow(2,cycle_depth)+1;j++) {
                sum+=pow(u_func(x[j],y[i])-u[i*(n*pow(2,cycle_depth)+1)+j],2.0);
            }
        }
        sum=pow(sum/(n*pow(2,cycle_depth+1)*n*pow(2,cycle_depth+1)),0.5);
        std::cout<<"test"<<j<<"'s "<<"absolute error in 2-norm: "<<sum<<" "<<std::endl;
    }
    void relative_error_2norm_print(int j) {
        double sum1=0.0,sum2=0.0;
        for(int i=0;i<n*pow(2,cycle_depth)+1;i++) {
            for(int j=0;j<n*pow(2,cycle_depth)+1;j++) {
                sum1+=pow(u_func(x[j],y[i])-u[i*(n*pow(2,cycle_depth)+1)+j],2.0);
                sum2+=pow(u_func(x[j],y[i]),2.0);
            }
        }
        sum1=pow(sum1/(n*pow(2,cycle_depth)*n*pow(2,cycle_depth)),0.5);
        sum2=pow(sum2/(n*pow(2,cycle_depth)*n*pow(2,cycle_depth)),0.5);
        std::cout<<"test"<<j<<"'s "<<"relative error in 2-norm: "<<sum1/sum2<<" "<<std::endl;
    }
};

#else
#endif
